Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

RFC: Make it possible to register a LLVM type constructor for codegen. #50607

Draft
wants to merge 1 commit into
base: master
Choose a base branch
from
Draft

RFC: Make it possible to register a LLVM type constructor for codegen. #50607

wants to merge 1 commit into from
+35 −14

Conversation

maleadt
Copy link
Member

@maleadt maleadt commented Jul 20, 2023

This PR extends DataType with a field for an LLVM type constructor, so that its possible to generate code that uses unsupported LLVM types. The goal here is to extend BFloat16s.jl so that it can use LLVM's bfloat support, without specifically having to add bfloat16 support to Base. Demo:

primitive type BFloat16 16 end

# how to represent this type in LLVM
using LLVM
function bfloat16_type_ctor(ctx_ref)
    ctx = LLVM.Context(ctx_ref)
    LLVM.API.LLVMBFloatTypeInContext(ctx)
end
const bfloat16_type_cb =
    @cfunction(bfloat16_type_ctor, LLVM.API.LLVMTypeRef, (LLVM.API.LLVMContextRef,))
Core._typector!(BFloat16, bfloat16_type_cb)

# some simple functionality
Base.:(+)(x::BFloat16, y::BFloat16) = Base.llvmcall(("""
        define bfloat @add(bfloat %x, bfloat %y) #0 {
            %z = fadd bfloat %x, %y
            ret bfloat %z
        }

        attributes #0 = { alwaysinline }""", "add"),
    BFloat16, Tuple{BFloat16, BFloat16}, x, y)

# generate code
using InteractiveUtils
code_llvm(+, (BFloat16, BFloat16))
;  @ /Users/tim/Julia/src/julia/build/dev/wip.jl:14 within `+`
define bfloat @"julia_+_119"(bfloat %"x::BFloat16", bfloat %"y::BFloat16") #0 {
top:
  %z.i = fadd bfloat %"x::BFloat16", %"y::BFloat16"
  ret bfloat %z.i
}

One open question is how to handle native code generation. Most platforms don't support bfloat16, so LLVM will fail during instruction selection. Legalization is relatively straightforward, so we could have a demote pass like we have for Float16. However, do we also want that pass to live in BFloat16s.jl? It would be weird to sell this as a mechanism to support arbitrary types, while still requiring legalization passes in Julia's compiler...

Ref #41075

@maleadt maleadt added speculative Whether the change will be implemented is speculative compiler:codegen Generation of LLVM IR and native code labels Jul 20, 2023
@maleadt maleadt force-pushed the tb/llvm_type_ctor branch from be079cb to 9b51354 Compare July 20, 2023 13:26
@vchuravy
Copy link
Member

It would be weird to sell this as a mechanism to support arbitrary types, while still requiring legalization passes in Julia's compiler...

Yeah... I think I would prefer the type definition and the legalization to live in the Julia compiler itself, but I also understand that as a language we might now want to take the extra burden of doing soft-fp on every new ML type xD.

I am not sure I am a 100% convinced that the LLVM ctor is the best idea, since at least on paper we pretend that one could write a Julia compiler without using LLVM.

@maleadt
Copy link
Member Author

maleadt commented Jul 20, 2023

I am not sure I am a 100% convinced that the LLVM ctor is the best idea, since at least on paper we pretend that one could write a Julia compiler without using LLVM.

The problem then isn't necessarily with the LLVM type ctor, which I guess just wouldn't be used when doing interpretation, but with the llvmcalls used to implement BFloat16 functionality. Maybe we need something a la if @generated to provide runtime intrinsics:

function Base.:(+)(x::BFloat16, y::BFloat16)
    if @compiled
        Base.llvmcall("fadd bfloat16")
    else
        xbits = reinterpret(UInt16, x)
        ybits = reinterpret(UInt16, y)
        # bit twiddling
        reinterpret(BFloat16, ...)
    end
end

@Keno
Copy link
Member

Keno commented Jul 20, 2023

Can we just specify the LLVM type as a symbol and then parse it using whatever LLVM version you happen to be using? The semantics could also be such that if the type is not legal, we refuse to use it and just do the standard integer lowering.

@maleadt
Copy link
Member Author

maleadt commented Jul 20, 2023

The semantics could also be such that if the type is not legal, we refuse to use it and just do the standard integer lowering.

How do you determine if it's legal?

@Keno
Copy link
Member

Keno commented Jul 20, 2023

How do you determine if it's legal?

I figured it'd be a codegen option. That said, I also don't really have a problem with putting the legalization pass into base. The set of legal llvm types is finite and small. It seems unnecessary to design a super general extension mechanism for it. Of course, there are other backends we may want to consider in the future (XLA, MLIR, etc) and it'd be nice for this kind of mechanism to be flexible enough to adapt.

@gbaraldi
Copy link
Member

I guess a mechanism similar/the same as this would be interesting also for #45486

@maleadt maleadt mentioned this pull request Sep 26, 2023
Merged
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers
No reviews
Assignees
No one assigned
Labels
compiler:codegen Generation of LLVM IR and native code speculative Whether the change will be implemented is speculative
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.
4 participants
@maleadt @vchuravy @Keno @gbaraldi